//AnimatedModel.h
//Created 06/01/14
//Created By Daniel Bowler
//
//Loads and manages models with animation data. We supporty key frame animation. 
//
//Note: We support up to 256 bones in a single skeleton. Each vertex will be influenced
//by a maximum of 4 bones

#pragma once

//Need XNA and D3D11 libs
#include "d3dUtil.h"
#include "xnacollision.h"
//Include the math helper
#include "MathHelper.h"

//Include StaticModel so we have access to the TessellationData struct
#include "StaticModel.h"

//Includes aiQauternion etc. 
#include <assimp\anim.h>
//Include a math library that actually works
#include "math_3d.h"

//Maximum number of bones for a skeleton. 
#define MAX_NUMBER_OF_BONES_PER_SKELETON 256
#define MAX_BONES_PER_VERTEX 4

//Forward declare aiScene - data struct that contains all we need in order
//to load mesh, material and animation data
struct aiScene;
//aiMesh contains vertex data for a submodel. 
struct aiMesh;
//aiNode - node in the animation bone structure. May or may not match a bone, but
//will be used to generate final transformation matricies. 
struct aiNode;

class AnimatedModel
{
	//Friend declaration for the graphics interface allowing the graphics 
	//interface the option to access this classes private data - Used for
	//frustum culling amongst other things. 
	friend class GraphicsInterface;
public:
	//Standard constructor/Destrouctor
	AnimatedModel();
	~AnimatedModel();

	//Loads a mesh and animation data through Assimp - See StaticModel.h
	//for an explanation on the bool flags. Generally the default values will
	//work best, but some models are created in odd ways (or the artist is a moron...)
	void CreateAnimatedModelThroughAssimp(ID3D11Device* device, 
		std::string& file, std::string& textureDirectory, 		
		bool shouldMakeLeftHanded = true, bool shouldFlipUV = true, bool shouldFlipWinding = true,
		bool shouldFixInwardFacingNormals = false);

public:
	//Sets the world matrix data. 
	void SetWorldPosition(XMFLOAT3 pos);
	void SetWorldRotation(XMFLOAT3 rot);
	//Note: Make sure that the x, y and z scale factors are the same. Otherwise, 
	//frustum culling causes a crash. I may look in to fixing this at somepoint. 
	void SetWorldScale(XMFLOAT3 scale);
	void SetWorldMatrix(XMFLOAT3 pos, XMFLOAT3 rot, XMFLOAT3 scale);
	//Gets the world matrix data
	XMFLOAT3 GetWorldPosition(){return position;};
	XMFLOAT3 GetWorldRotation(){return rotation;};
	XMFLOAT3 GetWorldScale(){return scaling;};

	//Returns AABB in local space
	XNA::AxisAlignedBox GetAABBL(){return aabbL;};

	//Returns AABB in world space. 
	XNA::AxisAlignedBox CalculateWorldSpaceAABB();

	//Sets the shadow casting state. If true (default), then the model will
	//cast shadows. If false, the model will not. 
	void SetIsShadowCaster(bool state){isShadowCaster = state;};

	//Update Material Data.
	//
	//Firstly, updates all of the submodels materials with the same material
	void SetSubmodelsWithSameMaterial(Material& material);
	//Secondly, sets indivdual submodel (index) material.
	void SetSubModelMaterial(unsigned index, Material& material);
	//Thirdly, set reflection property for all submodels.
	void SetSubmodelsWithSameReflectionData(XMFLOAT3 &reflectivity);

public:
	//Sets the current active animation - If this animation doesnt exist, will
	//enter infinite loop. 
	//
	//These functions will reset the animationRunTimer variable to 0.0. Thus you
	//can use this is to reset the animation back to the start.
	void SetActiveAnimation(std::string& animationName);
	void SetActiveAnimation(unsigned animationIndex);
	//Returns the index of the current active animation. 
	unsigned GetActiveAnimationIndex(){return currentActiveAnimation;};

public:
	//Updates the animating model. This function will handle all the interpolation
	//required to generate the final transformation matricies. Simply pass the delta
	//value that is passed to the CurrentScene::Update() function along and eveything
	//should work fine (hopefully - although it is possible that the scaling
	//key frame data could cock things up).
	//
	//If you want a slower/faster animation, modify the passed in delta value. 
	//eg: animatedModel->Update(delta*2.0f); will see the animation race through
	//faster than the original modeler wanted. 
	//
	//delta passed to the current scene will never be a negative number. Though you
	//could pass a negative delta to the animation to have the animation reverse. However, 
	//this will cock things up once the animationRunTimer drops below 0 - my suggestion
	//is not to do it. 
	//
	//the shouldLoop flag, when true, will loop the animation indefinatly. When false, 
	//as soon as the animation completes, the animation will stop in its final pose. 
	//
	//NOTE: Make sure to set the current animation before calling this function!
	//
	//If riggedPose is true, draws the model in its rigged pose. 
	void Update(float delta, bool shouldLoop = true, bool riggedPose = false);

private:
	//Internal structure that represents a sub portion of a model (can also be the entire model - 
	//especially in the case of the simple geometry like boxes, grids etc)
	struct SubModelData
	{
		SubModelData()
		{
			vb = 0;
			ib = 0;
			TextureMap = 0;
			NormalMap = 0;
			HeightMap = 0;

			//Default topology is triangle list (basic models use this, some
			//.obj may use quads
			Topology = D3D11_PRIMITIVE_TOPOLOGY_TRIANGLELIST;

			StartDrawIndex = 0;
			DrawCount = 0;

			//Opaque objects.
			IsOpaque = true;

			//Initially not using displacement mapping.
			IsUsingDisplacementMapping = false;
			IsUsingSeperateHeightMap = false;

			//Init tex transform matrix.
			XMMATRIX i = XMMatrixIdentity();
			XMStoreFloat4x4(&TexTrans, i);
		}
		~SubModelData()
		{
			ReleaseCOM(vb);
			ReleaseCOM(ib);
			ReleaseCOM(TextureMap);
			ReleaseCOM(NormalMap);
			ReleaseCOM(HeightMap);
		}

		//Rendering data: Vertex buffer and optional index buffer. Per material
		//vb and ib (it is possible for multiple materials to share the same
		//vb / ib... but just incase the model fails to provide say tex coords for
		//a subsection of the model, the option to use individual vb/ib has been given)
		ID3D11Buffer* vb;
		ID3D11Buffer* ib;

		//Default topology is triangle list (basic models use this, some
		//.obj may use quads
		D3D11_PRIMITIVE_TOPOLOGY Topology;

		//Texture map
		ID3D11ShaderResourceView* TextureMap;
		//Normal map
		ID3D11ShaderResourceView* NormalMap;
		//Optional disp map
		ID3D11ShaderResourceView* HeightMap;
		
		//Flag indicating if this material uses displacement mapping. Makes
		//use of the normal map for this feature.
		bool IsUsingDisplacementMapping;
		bool IsUsingSeperateHeightMap;

		//Tessellation data for use with displacement mapping. 
		TessellationData TessData;

		//Lighting material
		Material Mat;

		//Number of indicies or verticies for this portion of the
		//model
		UINT StartDrawIndex;
		UINT DrawCount;

		//Bool indcating if this object is trasnaprant or opaque. IMPORTANT
		//otherwise lighting will be all messed up.
		bool IsOpaque;

		//Tex transformation matrix.
		XMFLOAT4X4 TexTrans;
	};

	//Animation data structures:
	//
	//Represents a bone in the skeleton. We will generate a final transform
	//matrix per bone. 
	struct Bone
	{
		Bone()
		{
			//Init matrix
			OffsetMatrix.InitIdentity();
		}
		std::string Name; //Name - identifies this bone - Must be unique.
		unsigned BoneIndex; //Bone ID. [0-bones.size()-1]
		Matrix4f OffsetMatrix; //Offset matrix for this bone. Needed for the final transforms. 
		unsigned BoneNodeIndex; //Index in to boneHierarchy. 
	};

	//Represents a bone node.
	struct BoneNode
	{
		BoneNode()
		{
			ParentIndex = 0;
			ToParentMatrix.InitIdentity();
			ChildNodes.clear();
		}
		~BoneNode()
		{
			ChildNodes.clear();
		}

		std::string Name; //Name of the node - if this node corrosponds to a bone, this will match a Bone::Name
		int BoneIndex;    // -1 if this node doesnt corrospond to a node. 
		Matrix4f ToParentMatrix; //To Parent NODE Matrix. 
		int ParentIndex; //Parent NODE (not bone). -1 is the root node.
		std::vector<BoneNode*>ChildNodes; //Child nodes.
	};

	//Filled with the final transforms in Mateix4f form. These will be
	//converted to XNA maths after gening the final transforms. 
	struct BoneInfo
    {
        Matrix4f FinalTransformation;        
    };

	//Key frame data - Each bone can have A number of translations, B number 
	//of rotations (Quaternions) and C number of scales. Each of these structs
	//represents one piece of keyframe data.
	struct KeyFrameTranslation
	{
		//Time at which the bone node is fully at this location
		double Time;
		//Translation vector
		aiVector3D Value;
	};

	struct KeyFrameRotation
	{
		//Time at which the bone node is fully at this location
		double Time;
		//Translation quaternion
		aiQuaternion Value;
	};

	struct KeyFrameScale
	{
		//Time at which the bone node is fully at this location
		double Time;
		//Translation quaternion
		aiVector3D Value;
	};

	//Represents a single bones animation data (orinentation). A bone can
	//have n number of KeyFrames. 
	struct BoneAnimation
	{
		//Destructor clears vectors. 
		~BoneAnimation()
		{
			Trans.clear();
			Rots.clear();
			Scales.clear();
		}

		//Each bone can have any number of rotations, translations and scales. 
		std::vector<KeyFrameTranslation>Trans;
		std::vector<KeyFrameRotation>Rots;
		std::vector<KeyFrameScale>Scales;
	};

	//Represents an animation clip - this could be "Walk", "Running" etc. Each animation
	//clip will contain an array of BoneAnimations - one per Bone in the skeleton. Each
	//BoneAnimation entry can contain any number of key frames. 
	struct AnimationClip
	{
		AnimationClip()
		{
			ClipName = "";
			BonesAnimations.clear();
		}
		~AnimationClip()
		{
			BonesAnimations.clear();
		}

		std::string ClipName; //Used to set/update current animation 
		double DurationTicks; //How long (in ticks) this animation runs for.
		double TickRate;      //How many ticks per second. 
		double AnimationTime; //How long this animation runs in seconds. 

		//An anmation clip will contain bonesHierarchy.size() BoneAnimation entires.Each 
		//of these represents one bone node in the skeleton. Note, some nodes
		//may not have any key frame data... but we still allocate enough
		//memory so that the ith entry in this vector corrosponds to the ith
		//entry in boneHierarchy vector. 
		std::vector<BoneAnimation>BonesAnimations;
	};

private:
	//System copy of the vertex positions of the mesh. Used to create an AABB
	//for frustum culling purposes. 
	std::vector<XMFLOAT3>systemVertexCopyArr; 
	//Axis-Aligned Boundign Box - for use with frusutm culling. NOT suitable for 
	//collision since the box has been scalled up. 
	XNA::AxisAlignedBox aabbL;

	//Vector of submodels - drawn individually. 
	std::vector<SubModelData*>subModels;

	//Animation Data
	//
	//Vector of animation clips - one for each animation this model loads. Eg, 
	//walk, run, idle etc. Default will be the animation at index 0
	std::vector<AnimationClip>animationClips;
	//Vector of Bones. This represents our skeleton mesh. Each bone will
	//have a final transformation matrix generated when we update the model.
	std::vector<Bone>bones;
	//Vector of BoneNodes - Essentially, our bone (node) hierarchy (tree). There
	//can be entires here that do not match a bone, but are required when we 
	//generate our bones final transforms. 
	//
	//boneHierarchy[0] will be the root NODE. 
	std::vector<BoneNode>boneHierarchy;
	//Root nodes global transformation matrix - inversed. 
	Matrix4f globalInverseTransform;

	//Array of final transforms - this is what we pass to the shaders. ith entry
	//corrosponds to the ith bone (not node). We will use an array over a vector since
	//we want continuouse memory. 
	//
	//Pre XNA conversion (Uses Matrix4f)
	std::vector<BoneInfo> boneInfo;
	//Post XNA conversion - This is what gets passed to the shaders. 
	XMFLOAT4X4* finalTransformArr;

private:
	//Flag set to true once we have loaded a model. Like the static model
	//classes, once an instance has been loaded with mesh data, we can not
	//recall CreateAnimatedModelThroughAssimp() again to generate new data.
	bool isInit;

	//Will this model cast shadows?
	bool isShadowCaster;

	//Total draw count - total number of indicies or verticies to draw all meterials. 
	//Maybe useful for the ZPrePass to cut down on Draw() calls (IE, objects that 
	//have several materials will be rendered in one z pre pass stage). 
	UINT totalDrawCount;

	//Global world matrix - essentially moves and orientates all verticies of a mesh.
	XMFLOAT4X4 worldMatrix;
	//Translation, rotation and scaling factors that make up the world matrix.
	XMFLOAT3 position;
	XMFLOAT3 rotation;
	XMFLOAT3 scaling;

	//Currently active animation. Defualts to the aimation at animationClips[0]; 
	unsigned currentActiveAnimation;

	//Total runtime of this animation. We += delta when the scene programmer calls the
	//update function
	double totalRunTimeOfAnimation;

private:
	//Loads bones from the animation data
	void LoadBones(const aiScene* scene);
	//Loads animation clips. 
	void LoadAnimationClips(const aiScene* scene);
	//Traverses the bone hierarchy tree and extracts whatever data we need
	//from it. The ToParent index and the ToParent matrix - should be called
	//AFTER LoadBones().
	void LoadBoneHierarchy(const aiScene* scene, bool shouldOutputNodeHierarchy = false);
	//Recursivly traverse the node tree to extract parent index and to parent matrix. 
	void LoadBoneChild(aiNode* child);

	//Loads a submodel through assimp. Once done, puts (back) the submodel
	//on to our vector, adds verticies to the system vertex copy, and updates totalDrawCount
	void LoadSubModelThroughAssimp(ID3D11Device* device, const aiScene* scene, unsigned meshIndex, 
		std::string& textureDir, bool loadDispMap);
	//Loads material data for a mesh (Sub Model)
	void LoadAssimpMaterial(ID3D11Device* device, aiMaterial* material, SubModelData* subModel,
		std::string &textureDir, bool loadDispMap);
	//Loads vertex data for a mesh (Sub Model) - also index data. Note, also fills the 
	//system copy of positions for use when we create our AABB. 
	void LoadAssimpVertexData(ID3D11Device* device, aiMesh* mesh, SubModelData* subModel);

private:
	//Converts from a string to wisde string. We need the wide string when loading a texture
	std::wstring ToWS(std::string s);
	//Uses D3DX to load a texture - png, jpeg, dds and (I think?) gif. Note, function
	//will detect if you are attempting to load a TGA. If so, it will load a file
	//by the same name but with png extention. Crashes if one doesnt exist. 
	void LoadTexture(ID3D11Device* device, std::string& file, ID3D11ShaderResourceView** srv);

	//Loops through our bones vector and ensures that the bone we are adding
	//to the vector hasnt already been added. 
	bool IsBoneUnique(Bone* b);
	//Returns the index of the entry in bones which matches the strng 
	//passed in. -1 if nothing found
	int ReturnIndexOfBone(std::string name);
	//Loops through our bones node vector and ensures the node were adding 
	//to the vector is unique. 
	bool IsBoneNodeUnique(BoneNode* b);
	//Returns the index of the entry in boneHierarchy which matches
	//the string passed in. -1 if nothing found. 
	int ReturnIndexOfBoneNode(std::string name);

private:
	//Creates world matrix.
	void CreateWorldMatrix();
	//Create AABB surounding the loaded mesh. 
	void CreateAABB();

private:
	//Called in the Update function. This function generates our final transform
	//matricies for use in the animation vertex shaders. 
	void GenerateFinalTransforms(bool riggedPose);

	//Called in GenFinalTransforms(), this function traverses the tree
	//structure and generates the final transform for a node. If this node
	//is a bone, then it will add the final transform to the finalTransforms
	//array of matricies. 
	//
	//Recusive function. Each node calls this function for each child node of its. 
	void TraverseNodes(float animationTime, 
		BoneNode* node, unsigned nodeIndex, const Matrix4f &toParentTransform);

	//Returns the interpolated scaling/translation vector or
	//rotation quaterion at this animation time for this bone node
	//index
	void CalcInterpolatedScaling(aiVector3D& out, float animationTime, unsigned boneNodeIndex);
    void CalcInterpolatedRotation(aiQuaternion& out, float animationTime, unsigned boneNodeIndex);
    void CalcInterpolatedPosition(aiVector3D& out, float animationTime, unsigned boneNodeIndex);    
   
	//Returns the index in to animationClips[].BoneAninimations[boneIndex].Trans/Rots/Scales
	//that we should use (and index+1) as our two vectors (or quaternions in the case
	//of rotations) to interpolate (linear for scaling and positions (/translations) and sphereical
	//for rotations)
	UINT FindScaling(float animationTime, UINT boneNodeIndex);
    UINT FindRotation(float animationTime, UINT boneNodeIndex);
    UINT FindPosition(float animationTime, UINT boneNodeIndex);

private:
	//Getters for the GraphicsInterface. Note that the GraphicsInterface is
	//a friend so it can access the private data/functions.
	//
	//Gets the final transformation matricies and the number of them. They will
	//be uploaded to the shaders for the vertex shader to use. 
	XMFLOAT4X4* GetFinalTransforms(){return &finalTransformArr[0];};
	unsigned GetNumFinalTransforms(){return bones.size();};
};